import numpy as np
import cv2

src=cv2.imread("/home/zj/Robot_Team/Homework/test.jpeg",cv2.IMREAD_COLOR)

kernel=cv2.getStructuringElement(cv2.MORPH_ELLIPSE,(9,9))

dst=cv2.dilate(src,kernel)
dst=cv2.erode(dst,kernel)

threshold_low=230
threshold_high=255

gray=cv2.cvtColor(dst,cv2.COLOR_BGR2GRAY)

ret,binary = cv2.threshold(gray,threshold_low,threshold_high,cv2.THRESH_BINARY)
binary=cv2.Canny(binary,threshold_low,threshold_high)

contours,hie=cv2.findContours(binary,cv2.RETR_TREE,cv2.CHAIN_APPROX_NONE)
cv2.drawContours(src,contours,-1,(0,0,255),2)

font=cv2.FONT_HERSHEY_SIMPLEX
circles = cv2.HoughCircles(binary,cv2.HOUGH_GRADIENT,1,100,param1=220,param2=30,minRadius=2,maxRadius=300)
circle = np.uint16(np.around(circles))

for cnt in contours:
    x,y,w,h=cv2.boundingRect(cnt)
    len_contours=cv2.arcLength(cnt,True)
    rate=len_contours/(2*(w+h))
    cv2.rectangle(src, (x, y), (x + w, y + h), (0, 0, 0), 2)
    if 0.9<=rate:
        cv2.putText(src,"Square",(x,y-5), font, 0.5,(0,0,255),1)
    else:
        if abs(circle[0,:][0][0]-(x+w/2))<50 and abs(370-(y+h/2))>100:
            cv2.putText(src,"Circle",(x,y-5), font, 0.5,(0,0,255),1)
        else:
            cv2.putText(src,"Triangle",(x,y-5), font, 0.5,(0,0,255),1)



cv2.imshow("Tackle",src)
cv2.waitKey(0)
cv2.destroyAllWindows()
